Fix JAX compatibility issues in Job Search III lecture #687

jstac · 2025-11-10T10:29:09Z

Summary

This PR fixes several JAX-related compatibility issues in the Job Search III: Search with Separation and Markov Wages lecture (mccall_model_with_sep_markov.md).

Changes Made

Main Fixes

Refactored vfi() function: Now returns only v_final instead of a tuple (v_star, σ_star), making the implementation cleaner and more consistent with typical VFI patterns
Removed successive_approx() function: Integrated the iteration logic directly into vfi() to simplify the codebase
Fixed get_reservation_wage() function: Rewrote to use jnp.argmax() instead of jnp.where() to avoid JAX concretization errors that occur during JIT compilation
Fixed JAX decorator imports:
- Changed @jit to @jax.jit
- Changed @partial(jit, ...) to @partial(jax.jit, ...)
- Removed unnecessary @jit decorators from T() and get_greedy() functions
Updated all vfi() call sites: Added explicit σ_star = get_greedy(v_star, model) calls after each vfi() invocation (7 locations total)

Additional Improvement

Minor wording improvement in mccall_model_with_separation.md for better clarity

Testing

✅ Converted the markdown file to Python using jupytext
✅ Ran the Python script successfully without errors
✅ Verified all MyST markdown structure is valid
✅ Confirmed jupytext can parse and convert the file

Technical Details

The main issue was that jnp.where() with dynamic size cannot be used inside JIT-compiled functions, causing a ConcretizationTypeError. The fix uses jnp.argmax() to find the first acceptance index, which works correctly with JAX's tracing mechanism.

🤖 Generated with Claude Code

@jit

Updated mccall_model_with_sep_markov.md to fix several JAX-related issues: - Refactored vfi() to return only v_final instead of tuple, making it more consistent with VFI pattern - Removed separate successive_approx() function and integrated iteration logic directly into vfi() - Fixed JAX decorators: changed @jit to @jax.jit and @partial(jit, ...) to @partial(jax.jit, ...) - Rewrote get_reservation_wage() to use jnp.argmax() instead of jnp.where() to avoid JAX concretization errors in JIT compilation - Updated all vfi() call sites to explicitly compute policy with get_greedy(v_star, model) - Removed @jit decorators from T() and get_greedy() functions (not needed) Also improved wording in mccall_model_with_separation.md for clarity. Tested: Converted to Python and ran successfully without errors. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

jstac · 2025-11-10T10:29:43Z

Key Fix: get_reservation_wage() Function

The most critical fix in this PR addresses a JAX concretization error. The original implementation used:

accept_indices = jnp.where(σ == 1)[0]
if len(accept_indices) == 0:
    return jnp.inf
return w_vals[accept_indices[0]]

This fails in JIT-compiled functions because:

jnp.where() requires a statically-known size when used with JAX tracing
The if len(...) check tries to concretize an abstract value

The new implementation uses:

first_accept_idx = jnp.argmax(σ)
return jnp.where(jnp.any(σ), w_vals[first_accept_idx], jnp.inf)

This works because:

jnp.argmax() finds the first True value without requiring dynamic sizing
jnp.where() is used with scalar condition (no dynamic size issues)
Everything remains traceable through JAX's JIT compilation

jstac · 2025-11-10T10:30:04Z

Refactoring Notes: vfi() Function

The vfi() function was simplified:

Before:

Used a separate successive_approx() helper function
Returned both value function and policy: return v_star, σ_star

After:

Iteration logic integrated directly into vfi()
Returns only value function: return v_final
Policy computed separately via σ_star = get_greedy(v_star, model)

Benefits:

Cleaner separation of concerns: VFI computes values, policy extraction is separate
More flexible: Users can compute value function without necessarily needing the policy
Consistent with typical VFI patterns: Most VFI implementations return just the value function
Less code: Removed 35 lines of the successive_approx() function

All 7 call sites in the lecture were updated to explicitly call get_greedy() after vfi().

jstac · 2025-11-10T10:30:22Z

Testing Details

The changes were thoroughly tested:

Conversion test: Used jupytext --to py to convert the markdown to Python

Execution test: Ran the Python script successfully with output:

Time-average unemployment rate (single agent): 0.2215
Cross-sectional unemployment rate (at t=200): 0.2257
Difference: 0.0042

Round-trip test: Converted back to markdown using jupytext --to myst
Validation checks:
- ✅ 21 code cells all properly formatted
- ✅ All MyST directives balanced (exercise/solution blocks)
- ✅ Math delimiters balanced (10 $$ blocks)
- ✅ YAML frontmatter valid
- ✅ Jupytext can parse without errors

All fixes are backwards-compatible in terms of functionality - the lecture produces the same numerical results and plots as before.

github-actions · 2025-11-10T10:45:57Z

📖 Netlify Preview Ready!

Preview URL: https://pr-687--sunny-cactus-210e3e.netlify.app (6573f2a)

📚 Changed Lecture Pages: mccall_model_with_sep_markov, mccall_model_with_separation

Updated the McCall model with separation to use h = u(c) + β * sum_w v_u(w) q(w) as the continuation value, matching the notation from the basic McCall model lecture. This makes the progression between lectures more intuitive for readers: - Basic model: h = c + β * sum_w v*(w) q(w) - Separation model: h = u(c) + β * sum_w v_u(w) q(w) Key changes: - Replaced scalar d with h throughout mathematical derivations - Updated closed-form expression for v_e(w) to use h - Modified iteration algorithm to solve for h instead of d - Simplified Bellman equations using h notation - Updated all code functions (compute_v_e, update_h, solve_model) - Changed plots and comments to reference h This improves consistency across the job search lecture series and makes the mathematical structure clearer by explicitly showing the continuation value includes both current utility and discounted future value. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

jstac · 2025-11-10T20:17:43Z

Standardized continuation value notation

I've updated the Job Search II lecture (McCall model with separation) to use more consistent notation with the basic McCall model.

Changes made:

Changed continuation value from d to h: Now using h = u(c) + β * sum_w v_u(w) q(w) instead of d = sum_w v_u(w) q(w)
This mirrors the basic model's h = c + β * sum_w v*(w) q(w), making the progression between lectures more intuitive
The notation is cleaner since h now directly represents the value of rejecting an offer

Benefits:

Better consistency: Readers can more easily connect concepts between Job Search I and II
Clearer mathematics: The continuation value h explicitly shows both current utility u(c) and discounted future value
Simpler equations: Bellman equation becomes v_u(w) = max{v_e(w), h} instead of max{v_e(w), u(c) + β*d}

All code has been updated and tested successfully. The mathematical derivations and implementation now align better with the pedagogical flow of the lecture series.

The JAX library is already included in the base environment and doesn't need to be explicitly installed, which was causing build failures. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Removed 'torchaudio' from the installation command for PyTorch.

github-actions · 2025-11-11T01:52:07Z

📖 Netlify Preview Ready!

Preview URL: https://pr-687--sunny-cactus-210e3e.netlify.app (7f7d6a5)

📚 Changed Lecture Pages: mccall_model_with_sep_markov, mccall_model_with_separation

mmcky · 2025-11-11T03:13:10Z

@jstac this quick fix looks good to go. All execution is passing and the two new lectures are building plots etc. with no Warnings.

jstac · 2025-11-11T03:21:11Z

Thanks @mmcky !! I'll check the lectures and merge if they are all good.

jstac · 2025-11-11T03:21:40Z

Are we still using pytorch btw @mmcky ?

ripgrep says no ...

jstac · 2025-11-11T03:22:09Z

Ah, sorry, I see you just removed it!

@mmcky

mmcky · 2025-11-11T03:23:04Z

Are we still using pytorch btw @mmcky ?

ripgrep says no ...

@jstac we use numpyro which uses pytorch (under the hood I believe) -- but it should install binaries.
I will run a full test build after merge for the cached lectures.

jstac · 2025-11-11T03:26:00Z

Thanks @mmcky

Another reason to break up these big lecture series, isn't it?

The issues in the build were nothing to do with my PRs. I had a pretty stressful morning trying to get the new lectures live and eventually giving up 😬

mmcky · 2025-11-11T03:28:10Z

Sorry for your morning @jstac. Infrastructure / install issues with JAX are stressful!

I am about to run a full test build on another branch -- but there must have been a nightly release of pytorch that has a pretty significant change in it.

jstac · 2025-11-11T05:48:54Z

Are we still using pytorch btw @mmcky ?
ripgrep says no ...

@jstac we use numpyro which uses pytorch (under the hood I believe) -- but it should install binaries. I will run a full test build after merge for the cached lectures.

@mmcky this is from perplexity:

Yes, you can use NumPyro without PyTorch. NumPyro is a lightweight probabilistic programming library that uses JAX as its backend for automatic differentiation and JIT compilation. It does not depend on PyTorch at all. In fact, NumPyro is designed to leverage JAX's NumPy-compatible API and supports running on CPU, GPU, or TPU through JAX.

mmcky · 2025-11-11T06:02:00Z

Sorry @jstac used the wrong name, should be pyro-ppl -- we install numpyro and pyro-ppl which was previously used. Just running a search to see when we removed it. I actually think @HumphreyYang removed the use of pyro-ppl fairly recently.

…gma to gamma and use glue for figures Updated the McCall model with separation lecture with the following changes: Key changes: - Changed utility function parameter from σ (sigma) to γ (gamma) - Moved γ default value from utility function to Model class (γ: float = 2.0) - Updated all functions (compute_v_e, update_h, solve_model) to pass γ parameter - Simplified model unpacking to use tuple unpacking directly (e.g., α, β, γ, c, w, q = model) - Replaced static PNG figures with myst-nb glue functionality - Added glue import and glue() calls in exercise solutions - Converted {figure} directives to {glue:figure} directives for dynamic figure generation Benefits: - More consistent parameter naming (gamma is standard for CRRA utility) - Better code organization with parameter defaults in Model class - Cleaner unpacking syntax - Dynamic figure generation eliminates need for static PNG files - Figures automatically stay in sync with code 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

jstac · 2025-11-11T07:09:24Z

Additional updates to `mccall_model_with_separation.md`

I've made several improvements to the McCall model with separation lecture:

Parameter naming

Changed utility function parameter from σ (sigma) to γ (gamma) for better consistency with standard CRRA utility notation
Moved the default value (γ = 2.0) from the utility function to the Model class definition

Code improvements

Simplified model parameter unpacking to use direct tuple unpacking: α, β, γ, c, w, q = model
Updated all functions (compute_v_e, update_h, solve_model) to properly pass the γ parameter

Dynamic figure generation with glue

Added myst_nb glue functionality to replace static PNG figures
Added glue() calls in exercise solution code to capture figures
Converted all {figure} directives to {glue:figure} directives
This eliminates the need for the three static PNG files and ensures figures are always in sync with the code

These changes improve code maintainability and follow better practices for Jupyter Book documentation.

github-actions · 2025-11-11T07:26:08Z

📖 Netlify Preview Ready!

Preview URL: https://pr-687--sunny-cactus-210e3e.netlify.app (deb99b1)

📚 Changed Lecture Pages: mccall_model_with_sep_markov, mccall_model_with_separation

Removed static PNG files that are now dynamically generated using myst-nb glue: - mccall_resw_alpha.png - mccall_resw_beta.png - mccall_resw_c.png These figures are now generated from the exercise solution code and displayed via glue:figure directives, eliminating the need for static files and ensuring figures always match the code. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

jstac · 2025-11-11T07:37:31Z

@mmcky @HumphreyYang

Me and my new best buddy Claude used glue to get rid of some static assets -- png files. Deleting files and dependencies is deeply satisfying. No more worrying about whether the png files are out of sync.

github-actions · 2025-11-11T07:54:18Z

📖 Netlify Preview Ready!

Preview URL: https://pr-687--sunny-cactus-210e3e.netlify.app (4e6859a)

📚 Changed Lecture Pages: mccall_model_with_sep_markov, mccall_model_with_separation

github-actions · 2025-11-11T08:01:13Z

📖 Netlify Preview Ready!

Preview URL: https://pr-687--sunny-cactus-210e3e.netlify.app (35aa16d)

📚 Changed Lecture Pages: mccall_model_with_sep_markov, mccall_model_with_separation

jstac and others added 7 commits November 11, 2025 05:28

Update PyTorch installation to use cu121

8ad69b1

Update CI workflow to install only PyTorch packages

66013bc

Removed 'torchaudio' from the installation command for PyTorch.

fix: JAX version pinning

702773e

install new CUDANN binaries

949043e

remove pytorch

7f7d6a5

mmcky mentioned this pull request Nov 11, 2025

TEST: full execution check #688

Closed

4 tasks

jstac and others added 2 commits November 11, 2025 16:09

misc

2ad5d43

misc

35aa16d

jstac merged commit 0d564d4 into main Nov 11, 2025
1 check passed

jstac deleted the js_iii_edits branch November 11, 2025 08:03

Uh oh!

Fix JAX compatibility issues in Job Search III lecture #687

Fix JAX compatibility issues in Job Search III lecture #687

Uh oh!

Conversation

jstac commented Nov 10, 2025

Summary

Changes Made

Main Fixes

Additional Improvement

Testing

Technical Details

Uh oh!

jstac commented Nov 10, 2025

Key Fix: get_reservation_wage() Function

Uh oh!

jstac commented Nov 10, 2025

Refactoring Notes: vfi() Function

Uh oh!

jstac commented Nov 10, 2025

Testing Details

Uh oh!

github-actions bot commented Nov 10, 2025

Uh oh!

jstac commented Nov 10, 2025

Standardized continuation value notation

Changes made:

Benefits:

Uh oh!

github-actions bot commented Nov 11, 2025

Uh oh!

mmcky commented Nov 11, 2025

Uh oh!

jstac commented Nov 11, 2025

Uh oh!

jstac commented Nov 11, 2025

Uh oh!

jstac commented Nov 11, 2025

Uh oh!

mmcky commented Nov 11, 2025

Uh oh!

jstac commented Nov 11, 2025

Uh oh!

mmcky commented Nov 11, 2025

Uh oh!

jstac commented Nov 11, 2025

Uh oh!

mmcky commented Nov 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jstac commented Nov 11, 2025

Additional updates to mccall_model_with_separation.md

Parameter naming

Code improvements

Dynamic figure generation with glue

Uh oh!

github-actions bot commented Nov 11, 2025

Uh oh!

jstac commented Nov 11, 2025

Uh oh!

github-actions bot commented Nov 11, 2025

Uh oh!

github-actions bot commented Nov 11, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

mmcky commented Nov 11, 2025 •

edited

Loading

Additional updates to `mccall_model_with_separation.md`